import torch
import torch.nn as nn
import torch.nn.functional as F


# 自定义神经网络类
class VGGBase(nn.Module):
    def __init__(self):
        super(VGGBase, self).__init__()
        # =====================================================================
        # 卷积层 3*28*28 -》 64*14*14
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )
        # 池化层
        self.max_pooling1 = nn.MaxPool2d(kernel_size=2, stride=2)

        # =====================================================================
        # 卷积层 64*14*14 -》 128*7*7
        self.conv2_1 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )
        self.conv2_2 = nn.Sequential(
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )
        # 池化层
        self.max_pooling2 = nn.MaxPool2d(kernel_size=2, stride=2)

        # =====================================================================
        # 卷积层 128*7*7 -》 256*4*4
        self.conv3_1 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU()
        )
        self.conv3_2 = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU()
        )
        # 池化层
        self.max_pooling3 = nn.MaxPool2d(kernel_size=2, stride=2, padding=1)

        # =====================================================================
        # 卷积层 256*4*4 -》 512*2*2
        self.conv4_1 = nn.Sequential(
            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU()
        )
        self.conv4_2 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU()
        )
        # 池化层
        self.max_pooling4 = nn.MaxPool2d(kernel_size=2, stride=2)
        # =====================================================================
        # fc全连接层 batchSize * 512*2*2 -> batchSize * (512*4) -> 10 labels 二维张量Tensor

        self.fc = nn.Linear(512 * 4, 10)

    def forward(self, x):
        # X IS batchsize * 3 * 28 * 28
        batchsize = x.size(0)
        out = self.conv1(x)
        out = self.max_pooling1(out)
        out = self.conv2_1(out)
        out = self.conv2_2(out)
        out = self.max_pooling2(out)
        out = self.conv3_1(out)
        out = self.conv3_2(out)
        out = self.max_pooling3(out)
        out = self.conv4_1(out)
        out = self.conv4_2(out)
        out = self.max_pooling4(out)
        out = out.view(batchsize, -1)
        out = self.fc(out)
        # batchsize *(514*4) -> batchsize * 10
        out = F.log_softmax(out, dim=1)
        return out


def VGGNet():
    return VGGBase()
